package com.smart.thrift.config.client;

import cn.hutool.core.util.ClassUtil;
import cn.hutool.core.util.ReflectUtil;
import com.smart.thrift.base.Invocation;
import com.smart.thrift.base.InvocationImpl;
import com.smart.thrift.base.ThriftInterceptor;
import com.smart.thrift.base.ThriftSerializer;
import com.smart.thrift.discoveryandbalance.balancer.AbstractServerListLoadBalancer;
import com.smart.thrift.discoveryandbalance.balancer.ILoadBalancer;
import com.smart.thrift.discoveryandbalance.discovery.ThriftServerNode;
import com.smart.thrift.discoveryandbalance.rule.IRule;
import com.smart.thrift.discoveryandbalance.rule.RoundRobinRule;
import com.smart.thrift.exception.client.ThriftApplicationException;
import com.smart.thrift.exception.client.ThriftClientException;
import com.smart.thrift.exception.client.ThriftClientOpenException;
import com.smart.thrift.exception.client.ThriftClientRequestTimeoutException;
import com.smart.thrift.model.common.ThriftClientProperties;
import com.smart.thrift.model.common.ThriftServiceSignature;
import com.smart.thrift.model.pool.TransportKeyedObjectPool;
import com.smart.thrift.thriftapi.BaseException;
import com.smart.thrift.thriftapi.BaseRequest;
import com.smart.thrift.thriftapi.BaseResponse;
import com.smart.thrift.thriftapi.BasicService;
import lombok.extern.slf4j.Slf4j;
import org.apache.thrift.TApplicationException;
import org.apache.thrift.TException;
import org.apache.thrift.TServiceClient;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TMultiplexedProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.springframework.context.ApplicationContext;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.UndeclaredThrowableException;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

@Slf4j
public class ThriftClientInvocationHandler implements InvocationHandler {

    private final ThriftServiceSignature serviceSignature;
    private final Class<?> clientClass;
    private final Constructor<? extends TServiceClient> clientConstructor;
    private final Class<?> beanClass;
    private final ApplicationContext applicationContext;

    private ILoadBalancer<?> loadBalancer;
    private ThriftClientProperties properties;
    private TransportKeyedObjectPool objectPool;
    private List<ThriftInterceptor> interceptors;
    private ThriftSerializer thriftSerializer;


    private final AtomicBoolean initialized = new AtomicBoolean(false);

    public ThriftClientInvocationHandler(ThriftServiceSignature serviceSignature, Class<?> clientClass,Class<?> beanClass,
                                         Constructor<? extends TServiceClient> clientConstructor, ApplicationContext applicationContext) throws Exception {
        this.serviceSignature = serviceSignature;
        this.clientClass = clientClass;
        this.beanClass = beanClass;
        this.clientConstructor = clientConstructor;
        this.applicationContext = applicationContext;
    }

    private void init() {
        this.thriftSerializer = applicationContext.getBean(ThriftSerializer.class);
        ThriftClientManager thriftClientManager = applicationContext.getBean(ThriftClientManager.class);
        this.properties = thriftClientManager.getProperties();
        this.objectPool = thriftClientManager.getObjectPool();
        this.interceptors = applicationContext.getBeansOfType(ThriftInterceptor.class).values().stream()
                .sorted((o1, o2) -> o1.getOrder().compareTo(o2.getOrder())).collect(Collectors.toList());
        IRule routerRule = new RoundRobinRule();
        this.loadBalancer = AbstractServerListLoadBalancer.getServerListLoadBalancer(routerRule);
        routerRule.setLoadBalancer(loadBalancer);
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] arguments) throws Throwable {
        if (!initialized.get()) {
            init();
        }
        switch (method.getName()) {
            case "toString":
                return beanClass.getName();
            case "equals":
                Object serviceSignature1;
                return arguments != null && arguments.length == 1 && ClassUtil.isAssignable(beanClass, arguments[0].getClass()) &&
                        (serviceSignature1 = ReflectUtil.getFieldValue(arguments[0], "serviceSignature")) != null && serviceSignature.equals(serviceSignature1);
            case "hashCode":
                return this.hashCode();
            default:
                ThriftClientProperties.ThriftClientPoolProperties poolProperties = properties.getPool();

                ThriftServerNode serverNode = null;
                String marker = serviceSignature.getServiceName();
                Set<String> markers = serviceSignature.markersByServiceClass();
                for (String m : markers) {
                    if (loadBalancer.containsServerNode(m)) {
                        marker = m;
                        serverNode = loadBalancer.chooseServerNode(m);
                        break;
                    }
                }

                if (Objects.isNull(serverNode)) {
                    return null;
                }

                int retryTimes = 0;
                TTransport transport = null;

                while (true) {
                    if (retryTimes++ > poolProperties.getRetryTimes()) {
                        log.error("All thrift client call failed, method is {}, args is {}, retryTimes: {}", method.getName(), arguments, retryTimes);
                        throw new ThriftClientException("Thrift client call failed, thrift client signature is: " + marker);
                    }

                    try {
                        transport = objectPool.borrowObject(serverNode);
                        TProtocol protocol = new TCompactProtocol(transport);
                        TMultiplexedProtocol multiplexedProtocol = new TMultiplexedProtocol(protocol, marker);
                        Object client = clientConstructor.newInstance(multiplexedProtocol);

                        if (ClassUtil.isAssignable(BasicService.Client.class, clientClass)) {
                            BasicService.Client client1 = (BasicService.Client) client;
                            BaseRequest baseRequest = new BaseRequest();
                            baseRequest.setMethodName(method.getName());
                            baseRequest.setParamTypes(Arrays.stream(arguments).map(arg -> arg.getClass().getName()).collect(Collectors.toList()));
                            baseRequest.setParams(Arrays.stream(arguments).map(thriftSerializer::serializeParameter).collect(Collectors.toList()));
                            Map<String, String> header = new HashMap<>();
                            if (!interceptors.isEmpty()) {
                                Invocation invocation1 = new InvocationImpl(serviceSignature.getServiceName(), method.getName(), serviceSignature.getTargetServiceClass().getName(), Arrays.stream(arguments).map(Object::getClass).collect(Collectors.toList()).toArray(new Class[arguments.length]), arguments, header);
                                boolean match = interceptors.stream().anyMatch(thriftInterceptor -> thriftInterceptor.preHandle(invocation1));
                                if (!match) {
                                    return null;
                                }
                            }
                            baseRequest.setHeader(header);
                            BaseResponse send = client1.send(baseRequest);
                            String dataType = send.getDataType();
                            Object result = null;
                            if (Objects.nonNull(dataType)) {
                                Class<?> resultType;
                                try {
                                    resultType = Class.forName(dataType);
                                } catch (ClassNotFoundException e) {
                                    Class<?> returnType = method.getReturnType();
                                    throw new ClassCastException("return type mismatch, Can't convert "+dataType+" to " + returnType.getName());
                                }
                                result = thriftSerializer.deserializeReturnValue(send.getData(), resultType);
                            }
                            for (ThriftInterceptor interceptor : interceptors) {
                                result = interceptor.postHandle(result);
                            }
                            return result;
                        } else {
                            Method cachedMethod = ReflectUtil.getMethod(client.getClass(), method.getName(), method.getParameterTypes());
                            return ReflectionUtils.invokeMethod(cachedMethod, client, arguments);
                        }
                    } catch (BaseException baseException) {
                        throw baseException;
                    } catch (IllegalArgumentException | IllegalAccessException | InstantiationException |
                             SecurityException | NoSuchMethodException e) {
                        throw new ThriftClientOpenException("Unable to open thrift client", e);
                    } catch (UndeclaredThrowableException e) {
                        Throwable undeclaredThrowable = e.getUndeclaredThrowable();
                        if (undeclaredThrowable instanceof TTransportException) {
                            TTransportException innerException = (TTransportException) e.getUndeclaredThrowable();
                            Throwable realException = innerException.getCause();

                            if (realException instanceof SocketTimeoutException) { // 超时,直接抛出异常,不进行重试
                                if (transport != null) {
                                    transport.close();
                                }

                                log.error("Thrift client request timeout, ip is {}, port is {}, timeout is {}, method is {}, args is {}",
                                        serverNode.getHost(), serverNode.getPort(), serverNode.getTimeout(),
                                        method, arguments);
                                throw new ThriftClientRequestTimeoutException("Thrift client request timeout", e);

                            } else if (realException == null && innerException.getType() == TTransportException.END_OF_FILE) {
                                // 服务端直接抛出了异常 or 服务端在被调用的过程中被关闭了
                                objectPool.clear(serverNode); // 把以前的对象池进行销毁
                                if (transport != null) {
                                    transport.close();
                                }

                            } else if (realException instanceof SocketException) {
                                objectPool.clear(serverNode);
                                if (transport != null) {
                                    transport.close();
                                }
                            }

                        } else if (undeclaredThrowable instanceof TApplicationException) {  // 有可能服务端返回的结果里存在null
                            log.error(
                                    "Thrift end of file, ip is {}, port is {}, timeout is {}, method is {}, args is {}, retryTimes is {}",
                                    serverNode.getHost(), serverNode.getPort(), serverNode.getTimeout(),
                                    method, arguments, retryTimes);
                            if (retryTimes >= poolProperties.getRetryTimes()) {
                                throw new ThriftApplicationException("Thrift end of file", e);
                            }

                            objectPool.clear(serverNode);
                            if (transport != null) {
                                transport.close();
                            }

                        } else if (undeclaredThrowable instanceof TException) { // idl exception
                            throw undeclaredThrowable;
                        } else {
                            // Unknown Exception
                            throw e;
                        }

                    } catch (Exception e) {
                        if (e instanceof ThriftClientOpenException) { // 创建连接失败
                            Throwable realCause = e.getCause().getCause();
                            // unreachable, reset router
                            if (realCause instanceof SocketException && realCause.getMessage().contains("Network is unreachable")) {
                                throw e;
                            }
                        } else {
                            throw e;
                        }
                    } finally {
                        try {
                            if (objectPool != null && transport != null) {
                                objectPool.returnObject(serverNode, transport);
                            }
                        } catch (Exception e) {
                            log.error(e.getMessage(), e);
                        }
                    }
                }
        }
    }
}
